# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： resize
Description :
Author : 'li'
date： 2022/1/27
-------------------------------------------------
Change Activity:
2022/1/27:
-------------------------------------------------
"""
import cv2

from ml.dl.misc.img_augment.resize_img import resize_img


class ResizeImg:
    def __init__(self, shape=None, scale=None, input_key='img', output_key='img', resize_target=False):
        """

        Args:
            shape:
            scale:
            input_key:
            output_key:
            resize_target:
        """

        self.shape = shape
        self.scale = scale
        self.input_key = input_key
        self.output_key = output_key
        self.resize_target = resize_target

    @staticmethod
    def _resize(img, shape):
        """

        Args:
            img:
            shape:

        Returns:

        """
        dtype = img.dtype
        # if dtype != np.uint8:
        #     img = img.astype(np.uint8)
        return cv2.resize(img, shape).astype(dtype)

    def __call__(self, data):
        img = data[self.input_key].copy()
        resized_image, resized_target = resize_img(img, self.shape, self.scale)
        if not self.resize_target:
            data[self.output_key] = resized_image
        else:
            data[self.output_key] = [resized_image, resized_target]
        return data
